from rembg import remove, new_session
from PIL import Image
import argparse
import os
import sys
import glob
import uuid  # 添加uuid模块用于生成唯一哈希值

# 导入颜色工具模块
from color_utils import RED, GREEN, YELLOW, BLUE, PURPLE, CYAN, BOLD, NC

def generate_unique_hash():
    """生成一个唯一的哈希值"""
    return uuid.uuid4().hex[:8]  # 使用8位哈希值，可以根据需要调整长度

def process_image(input_path, output_path=None, bg_color=None):
    """
    使用rembg库处理图片，移除背景
    
    Args:
        input_path: 输入图片路径
        output_path: 输出图片路径，如果为None则自动生成
        bg_color: 背景颜色，格式为十六进制颜色代码，如"#FF0000"
    
    Returns:
        bool: 处理是否成功
    """
    try:
        # 打开输入图片
        input_image = Image.open(input_path)
        
        # 如果没有指定输出路径，则根据输入文件名和哈希值生成
        if output_path is None:
            input_filename = os.path.basename(input_path)
            name, _ = os.path.splitext(input_filename)
            unique_hash = generate_unique_hash()
            # 如果有背景色，使用原始扩展名，否则使用PNG
            output_ext = '.png'
            output_path = os.path.join(os.path.dirname(input_path), f"{name}_{unique_hash}{output_ext}")
        
        # 使用rembg移除背景
        print(f"{CYAN}正在处理图片: {BLUE}{os.path.basename(input_path)}{NC}...{NC}")
        session = new_session(model_name="u2net")
        
        # 处理背景色
        bgcolor = None
        if bg_color:
            try:
                # 移除可能的#前缀
                if bg_color.startswith('#'):
                    bg_color = bg_color[1:]
                    
                # 解析颜色
                r = int(bg_color[0:2], 16)
                g = int(bg_color[2:4], 16)
                b = int(bg_color[4:6], 16)
                
                bgcolor = (r, g, b, 255)
                print(f"{CYAN}使用背景颜色: {PURPLE}RGB({r}, {g}, {b}){NC}")
            except Exception as e:
                print(f"{YELLOW}背景颜色格式错误: {str(e)}，将使用透明背景{NC}")
                bgcolor = None
        
        # 执行抠图
        output_image = remove(input_image, session=session, bgcolor=bgcolor)
        
        # 确保输出目录存在
        output_dir = os.path.dirname(output_path)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        # 保存输出图片
        # 如果是JPEG/JPG格式，需要将RGBA转换为RGB
        output_ext = os.path.splitext(output_path)[1].lower()
        if output_ext in ['.jpg', '.jpeg'] and output_image.mode == 'RGBA':
            # 创建RGB图像并粘贴RGBA图像
            rgb_image = Image.new('RGB', output_image.size, (255, 255, 255))
            rgb_image.paste(output_image, mask=output_image.split()[3])
            rgb_image.save(output_path, 'JPEG', quality=95)
        else:
            output_image.save(output_path)
        
        print(f"{GREEN}抠图成功: {BLUE}{os.path.basename(output_path)}{NC}")
        return True
    except Exception as e:
        print(f"{RED}处理图片时出错: {str(e)}{NC}")
        return False

def batch_process(input_dir, output_dir, bg_color=None):
    """
    批量处理目录下的所有图片
    
    Args:
        input_dir: 输入目录路径
        output_dir: 输出目录路径
        bg_color: 背景颜色
        
    Returns:
        bool: 是否全部处理成功
    """
    # 支持的图片格式
    image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.gif', '*.webp']
    
    # 获取所有图片文件
    image_files = []
    for ext in image_extensions:
        image_files.extend(glob.glob(os.path.join(input_dir, ext)))
        image_files.extend(glob.glob(os.path.join(input_dir, ext.upper())))
    
    if not image_files:
        print(f"{RED}错误: 在目录 {BLUE}{input_dir}{NC} 中未找到任何图片文件{NC}")
        return False
    
    print(f"{CYAN}找到 {BOLD}{len(image_files)}{NC}{CYAN} 个图片文件，开始批量处理...{NC}")
    
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 创建会话（只创建一次）
    session = new_session(model_name="u2net")
    
    # 处理背景色
    bgcolor = None
    if bg_color:
        try:
            # 移除可能的#前缀
            if bg_color.startswith('#'):
                bg_color = bg_color[1:]
                
            # 解析颜色
            r = int(bg_color[0:2], 16)
            g = int(bg_color[2:4], 16)
            b = int(bg_color[4:6], 16)
            
            bgcolor = (r, g, b, 255)
            print(f"{CYAN}使用背景颜色: {PURPLE}RGB({r}, {g}, {b}){NC}")
        except Exception as e:
            print(f"{YELLOW}背景颜色格式错误: {str(e)}，将使用透明背景{NC}")
            bgcolor = None
    
    # 处理每个图片
    success_count = 0
    for input_file in image_files:
        try:
            # 获取文件名和扩展名
            filename = os.path.basename(input_file)
            name, ext = os.path.splitext(filename)
            
            # 设置输出文件路径，添加唯一哈希值
            unique_hash = generate_unique_hash()
            output_ext = '.png' if bgcolor is None else ext
            output_file = os.path.join(output_dir, f"{name}_{unique_hash}{output_ext}")
            
            print(f"{CYAN}处理: {BLUE}{filename}{NC} -> {BLUE}{os.path.basename(output_file)}{NC}")
            
            # 打开输入图片
            input_image = Image.open(input_file)
            
            # 执行抠图
            output_image = remove(input_image, session=session, bgcolor=bgcolor)
            
            # 保存输出图片
            # 如果是JPEG/JPG格式，需要将RGBA转换为RGB
            if output_ext.lower() in ['.jpg', '.jpeg'] and output_image.mode == 'RGBA':
                # 创建RGB图像并粘贴RGBA图像
                rgb_image = Image.new('RGB', output_image.size, (255, 255, 255))
                rgb_image.paste(output_image, mask=output_image.split()[3])
                rgb_image.save(output_file, 'JPEG', quality=95)
            else:
                output_image.save(output_file)
            
            success_count += 1
        except Exception as e:
            print(f"{RED}处理 {BLUE}{input_file}{NC} 时出错: {str(e)}{NC}")
    
    print(f"{GREEN}批量处理完成: 成功 {BOLD}{success_count}/{len(image_files)}{NC}{GREEN} 个文件{NC}")
    return success_count > 0

def print_help():
    """打印帮助信息"""
    print(f"{CYAN}{BOLD}图片抠图工具 - 移除图片背景{NC}")
    print("")
    print(f"{BOLD}用法:{NC}")
    print(f"  {YELLOW}xzimg koutu -i <输入文件> -o <输出文件> [-c <背景色>]{NC}")
    print(f"  {YELLOW}xzimg koutu -b [输入目录] -o [输出目录] [-c <背景色>]{NC}")
    print("")
    print(f"{BOLD}参数:{NC}")
    print(f"  {GREEN}-i, --input{NC} <文件>    指定输入图片文件路径")
    print(f"  {GREEN}-o, --output{NC} <文件>   指定输出图片文件路径或目录")
    print(f"  {GREEN}-b, --batch{NC} [目录]    批量处理模式，可选指定输入目录，默认为当前目录")
    print(f"  {GREEN}-c, --color{NC} <颜色>    指定背景颜色，格式为十六进制，如 #FFFFFF")
    print(f"  {GREEN}--help{NC}               显示此帮助信息")
    print("")
    print(f"{BOLD}支持的图片格式:{NC}")
    print(f"  {PURPLE}png, jpg/jpeg, webp, gif, bmp, tiff/tif{NC}")
    print("")
    print(f"{BOLD}示例:{NC}")
    print(f"  {YELLOW}xzimg koutu -i input.jpg -o output.png{NC}")
    print(f"  {YELLOW}xzimg koutu -i input.jpg -o output.png -c \"#FF0000\"{NC}  # 使用红色背景")
    print(f"  {YELLOW}xzimg koutu -b -o output_dir{NC}                           # 批量处理当前目录下的所有图片")
    print(f"  {YELLOW}xzimg koutu -b images -o output_dir{NC}                    # 批量处理images目录下的所有图片")

def main():
    # 创建参数解析器
    parser = argparse.ArgumentParser(description='图片抠图工具 - 移除图片背景', add_help=False)
    
    # 添加帮助参数
    parser.add_argument('--help', action='store_true', help='显示帮助信息')
    
    # 创建互斥组，确保-i和-b不会同时使用
    group = parser.add_mutually_exclusive_group()
    group.add_argument('-i', '--input', help='输入图片路径')
    group.add_argument('-b', '--batch', nargs='?', const=os.getcwd(), help='批量处理的输入目录路径，不指定值则使用当前目录')
    
    parser.add_argument('-o', '--output', help='输出图片路径或目录')
    parser.add_argument('-c', '--color', help='背景颜色，格式为十六进制，如 #FFFFFF')
    
    # 解析命令行参数
    args, unknown = parser.parse_known_args()
    
    # 如果有--help参数，显示帮助信息并退出
    if args.help:
        print_help()
        return 0
    
    # 检查必要参数
    if not args.input and not args.batch:
        print(f"{RED}错误: 必须指定输入文件(-i)或批量处理模式(-b){NC}")
        print(f"使用 {YELLOW}--help{NC} 查看帮助信息")
        return 1
    
    # 判断是单文件模式还是批量模式
    if args.input:
        # 单文件模式
        if not os.path.isfile(args.input):
            print(f"{RED}错误: 输入文件不存在: {BLUE}{args.input}{NC}")
            return 1
        
        # 处理图片 - 如果没有指定输出路径，则传入None让函数自动生成
        success = process_image(args.input, args.output, args.color)
    else:
        # 批量模式
        input_dir = args.batch
        
        # 如果没有指定输出目录，默认为当前目录下的koutu文件夹
        if not args.output:
            output_dir = os.path.join(os.getcwd(), 'koutu')
        else:
            output_dir = args.output
        
        if not os.path.isdir(input_dir):
            print(f"{RED}错误: 输入目录不存在: {BLUE}{input_dir}{NC}")
            return 1
        
        # 批量处理图片
        success = batch_process(input_dir, output_dir, args.color)
    
    # 返回状态码
    return 0 if success else 1

if __name__ == "__main__":
    sys.exit(main())